import cv2
import numpy as np

KNOWN_WIDTH = 0.021  # 红色矩形的实际宽度为0.021米
FOCAL_LENGTH = 580  # 相机的焦距（单位：像素）


def find_marker(image):
    """
    检测图像中的红色矩形，并返回其轮廓
    """
    # 将图像从BGR转换为HSV颜色空间
    hsv = cv2.cvtColor(image, cv2.COLOR_BGR2HSV)

    # 定义红色的HSV范围
    lower_red1 = np.array([0, 120, 70])
    upper_red1 = np.array([10, 255, 255])
    lower_red2 = np.array([170, 120, 70])
    upper_red2 = np.array([180, 255, 255])

    # 创建红色的掩膜
    mask1 = cv2.inRange(hsv, lower_red1, upper_red1)
    mask2 = cv2.inRange(hsv, lower_red2, upper_red2)
    mask = cv2.bitwise_or(mask1, mask2)

    # 形态学操作，去除噪声
    kernel = np.ones((5, 5), np.uint8)
    mask = cv2.morphologyEx(mask, cv2.MORPH_OPEN, kernel)
    mask = cv2.morphologyEx(mask, cv2.MORPH_CLOSE, kernel)

    # 寻找轮廓
    contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)

    if len(contours) > 0:
        # 找到最大的轮廓
        c = max(contours, key=cv2.contourArea)
        return cv2.minAreaRect(c)
    else:
        return None


def distance_to_camera(knownWidth, focalLength, perWidth):
    """
    根据已知宽度、焦距和图像中的宽度计算距离
    """
    return (knownWidth * focalLength) / perWidth


# 初始化摄像头
cap = cv2.VideoCapture(0)

while True:
    # 读取一帧
    ret, frame = cap.read()
    if not ret:
        break

    # 检测红色矩形
    marker = find_marker(frame)

    if marker:
        # 获取矩形的宽度（单位：像素）
        perWidth = marker[1][0]

        # 计算距离
        distance = distance_to_camera(KNOWN_WIDTH, FOCAL_LENGTH, perWidth)

        # 绘制矩形和距离信息
        box = cv2.boxPoints(marker)
        box = np.int32(box)
        cv2.drawContours(frame, [box], -1, (0, 255, 0), 2)
        cv2.putText(frame, "Distance: %.2f m" % distance, (10, 30),
                    cv2.FONT_HERSHEY_SIMPLEX, 1.0, (0, 255, 0), 2)

    # 显示结果
    cv2.imshow("Frame", frame)

    # 按下'q'键退出
    if cv2.waitKey(1) & 0xFF == ord('q'):
        break

# 释放摄像头并关闭窗口
cap.release()
cv2.destroyAllWindows()
